# -*- encoding: utf-8 -*-
'''
@File    :   signal_processing.py
@Time    :   2022/3/22 14:18
@Author  :   ZhangChaoYang
@Desc    :   信号处理函数
'''

import copy

import pandas as pd
import numpy as np
from scipy.fft import rfft, irfft, rfftfreq
import pywt
from scipy.signal import stft, hilbert, chirp
from pyhht import EMD


def fft(y, fs):
    '''
    针对实数的快速傅利叶变换，把时域转到频域
    '''
    N = y.shape[-1]
    yf = rfft(y) / N  # 除以len(y)是为了把yf的值转到[0,1]上。len(yf)=len(y)/2+1
    xf = rfftfreq(N, 1 / fs)  # len(xf)=len(y)/2+1
    # print(xf.shape[-1], yf.shape[-1], N)
    return (xf, np.abs(yf))


def point_pass(y, fs, hz):
    '''
    只选取一个点的频率成分。点通滤波
    '''
    yy = rfft(y)  # 快速傅里叶变换，r是real（实数）的意思，只算一半，len(yy)是len(y)的一半
    N = len(y)
    xf = rfftfreq(N, 1 / fs)
    points_per_freq = len(xf) / (fs / 2)
    target_idx = int(hz * points_per_freq)
    # print(target_idx)
    yf = copy.copy(yy)
    yf[:target_idx - 1] = 0
    yf[target_idx + 1:] = 0
    return irfft(yf)


def band_pass(y, fs, lowHz, highHz):
    '''
    选取一个区间的频率成分。带通滤波
    '''
    yy = rfft(y)  # 快速傅里叶变换，r是real（实数）的意思，只算一半，len(yy)是len(y)的一半
    N = len(y)
    xf = rfftfreq(N, 1 / fs)
    points_per_freq = len(xf) / (fs / 2)
    low = int(lowHz * points_per_freq)
    high = int(highHz * points_per_freq)
    # print(low, high)
    yf = copy.copy(yy)
    yf[:low + 1] = 0
    yf[high + 1:] = 0
    return irfft(yf)


def multi_band_pass(y, fs, bands):
    '''
    选取一个区间的频率成分
    '''
    yy = rfft(y)
    N = len(y)
    xf = rfftfreq(N, 1 / fs)
    points_per_freq = len(xf) / (fs / 2)
    result = []
    for lowHz, highHz in bands:
        low = int(lowHz * points_per_freq)
        high = int(highHz * points_per_freq)
        yf = copy.copy(yy)
        yf[:low + 1] = 0
        yf[high + 1:] = 0
        result.append(irfft(yf))
    return result


def sfft(y, fs):
    '''
    快速傅利叶变换
    '''
    f, t, Zxx = stft(y, fs)
    magnitude = np.abs(Zxx)  # Zxx是复数，abs相当于求模长
    return (t, f, magnitude)


def cwt(y, fs, wavelet='cgau8'):
    '''
    连续小波变换
    通过print(pywt.wavelist())查看支持的wavelet：['bior1.1', 'bior1.3', 'bior1.5', 'bior2.2', 'bior2.4', 'bior2.6', 'bior2.8', 'bior3.1', 'bior3.3', 'bior3.5', 'bior3.7', 'bior3.9', 'bior4.4', 'bior5.5', 'bior6.8', 'cgau1', 'cgau2', 'cgau3', 'cgau4', 'cgau5', 'cgau6', 'cgau7', 'cgau8', 'cmor', 'coif1', 'coif2', 'coif3', 'coif4', 'coif5', 'coif6', 'coif7', 'coif8', 'coif9', 'coif10', 'coif11', 'coif12', 'coif13', 'coif14', 'coif15', 'coif16', 'coif17', 'db1', 'db2', 'db3', 'db4', 'db5', 'db6', 'db7', 'db8', 'db9', 'db10', 'db11', 'db12', 'db13', 'db14', 'db15', 'db16', 'db17', 'db18', 'db19', 'db20', 'db21', 'db22', 'db23', 'db24', 'db25', 'db26', 'db27', 'db28', 'db29', 'db30', 'db31', 'db32', 'db33', 'db34', 'db35', 'db36', 'db37', 'db38', 'dmey', 'fbsp', 'gaus1', 'gaus2', 'gaus3', 'gaus4', 'gaus5', 'gaus6', 'gaus7', 'gaus8', 'haar', 'mexh', 'morl', 'rbio1.1', 'rbio1.3', 'rbio1.5', 'rbio2.2', 'rbio2.4', 'rbio2.6', 'rbio2.8', 'rbio3.1', 'rbio3.3', 'rbio3.5', 'rbio3.7','rbio3.9', 'rbio4.4', 'rbio5.5', 'rbio6.8', 'shan', 'sym2', 'sym3', 'sym4', 'sym5', 'sym6', 'sym7', 'sym8', 'sym9', 'sym10', 'sym11', 'sym12', 'sym13', 'sym14', 'sym15', 'sym16', 'sym17','sym18', 'sym19', 'sym20']
    '''
    total_scale = 256
    fc = pywt.central_frequency(wavelet)
    cparm = 2 * fc * total_scale
    scales = cparm / np.arange(total_scale, 0, -1)
    cwtmatrix, frquency = pywt.cwt(y, scales, wavelet, 1.0 / fs)
    return frquency, np.abs(cwtmatrix)  # cwtmatrix是复数，abs相当于求模长


def ht(y, fs):
    '''
    希尔伯特变换
    '''
    analytic_signal = hilbert(y)  # 解调得到分析信号
    # amplitude_envelope = np.abs(analytic_signal)  # 求绝对值得到幅值包络
    instantaneous_phase = np.unwrap(np.angle(analytic_signal))  # 瞬时相位角
    instantaneous_frequency = (np.diff(instantaneous_phase) / (2.0 * np.pi) * fs)  # 瞬时频率
    return analytic_signal, instantaneous_frequency


def hht(y, fs):
    '''
    希尔伯特-黄 变换
    '''
    decomposer = EMD(y, n_imfs=5)  # 经验模态分解
    imfs = decomposer.decompose()  # 获取一系列固有模态函数
    results = [(y, ht(y, fs))]
    for imf in imfs:
        results.append((imf, ht(imf, fs)))
    return results


def hht_filter(y, fs, n_imfs, componentsRetain):
    '''
    抽取特定的固有模态函数
    '''
    decomposer = EMD(y, n_imfs=n_imfs)  # 经验模态分解
    imfs = decomposer.decompose()  # 获取一系列固有模态函数
    retain = np.sum(imfs[componentsRetain], axis=0)
    return retain


def statistics_feature(x):
    '''
    统计类特征
    @param x: 2-d arrray like or higher dimension
    '''
    data = pd.DataFrame(x)  # 从numpy转为DataFrame
    axis_num = len(data.shape)
    data_abs = data.abs()
    max = data.max(axis=axis_num - 1)  # 最大值
    min = data.min(axis=axis_num - 1)  # 最小值
    mean = data.mean(axis=axis_num - 1)  # 均值
    abs_mean = data_abs.mean(axis=axis_num - 1)
    var = data.var(axis=axis_num - 1)  # 方差
    std_var = data.std(axis=axis_num - 1)  # 标准差
    rms = np.sqrt(np.power(mean, 2) + var)  # 均方根
    skew = data.skew(axis=axis_num - 1)  # 偏度，为负时概率分布偏向于均值左侧
    kurt = data.kurt(axis=axis_num - 1)  # 峭度，概率分布的陡峭程度，又称为峰度
    form = rms / (data.mean(axis=axis_num - 1).abs())  # 波形因子
    peak = max / rms  # 峰值因子
    pulse = max / data_abs.mean(axis=axis_num - 1)  # 脉冲因子
    margin = max / np.power(np.sqrt(data_abs).sum(axis=axis_num - 1) / data.shape[-1], 2)  # 裕度
    square_sum = (data * data).sum(axis=axis_num - 1)  # 能量

    a = data.shift(periods=1, axis=axis_num - 1).drop(axis=axis_num - 1, labels=0)  # 右移一列
    a.columns -= 1
    b = data.shift(periods=-1, axis=axis_num - 1).drop(axis=axis_num - 1, labels=axis_num)  # 左移一列

    through_zero = (a * b < 0).sum(axis=axis_num - 1)  # 过零点数
    amplitude_diff = (a - b).abs().mean(axis=axis_num - 1)  # 平均幅值差
    aberrance = std_var / abs_mean  # 变异系数，反应数据的离散程度，由于除以了均值，它消除了量纲，具有很好的鲁棒性，不受到数据本身分布的影响
    return [min,  # 0 最小值
            max,  # 1 最大值
            mean,  # 2 平均值
            var,  # 3 方差
            std_var,  # 4 标准差
            rms,  # 5 有效值
            skew,  # 6 偏度
            kurt,  # 7 峭度
            form,  # 8 波形因子
            peak,  # 9 峰值因子
            pulse,  # 10 脉冲因子
            margin,  # 11 裕度
            square_sum,  # 12 能量
            through_zero,  # 13 过零点数
            amplitude_diff,  # 14 平均幅值差
            aberrance,  # 15 变异系数
            ]
